import numpy as np
import math
import sys
import pandas as pd
import os
import pymap3d as pm
import platform
import matplotlib.pyplot as plt
import json

import rosbag
from sensor_msgs.msg import NavSatFix, Imu
from nav_msgs.msg import Odometry
from nlink_parser.msg import LinktrackNodeframe3



def get_odomtry_data(t, msg):
    data = {
        # 'time':msg.header.stamp.secs + msg.header.stamp.nsecs*1e-9,
        'time': t.secs + t.nsecs*1e-9,
        'p_x': msg.pose.pose.position.x,
        'p_y': msg.pose.pose.position.y,
        'p_z': msg.pose.pose.position.z,
        'orient_x': msg.pose.pose.orientation.x,
        'orient_y': msg.pose.pose.orientation.y,
        'orient_z': msg.pose.pose.orientation.z,
        'orient_w': msg.pose.pose.orientation.w,
        'vel_linear_x': msg.twist.twist.linear.x,
        'vel_linear_y': msg.twist.twist.linear.y,
        'vel_linear_z': msg.twist.twist.linear.z,
        'vel_angular_x': msg.twist.twist.angular.x,
        'vel_angular_y': msg.twist.twist.angular.y,
        'vel_angular_z': msg.twist.twist.angular.z,
    }
    return data 

def get_nodeframe3_data(t, msg):
    """
    Extract data from a LinktrackNodeframe3 message.
    Handles nested LinktrackNode2 array within the message.
    """
    
    data = {
        'time': t.secs + t.nsecs*1e-9,
        'role': msg.role,
        'id': msg.id,
        'local_time': msg.local_time,
        'system_time': msg.system_time,
        'voltage': msg.voltage,
        'nodes': []
    }
    for node in msg.nodes:
        node_data = {
            'role': node.role,
            'id': node.id,
            'distance': node.dis,
            'fp_rssi': node.fp_rssi,
            'rx_rssi': node.rx_rssi,
        }
        data['nodes'].append(node_data)
    
    return data

def get_imu_data(t, msg):
    data = {
        # 'time':msg.header.stamp.secs + msg.header.stamp.nsecs*1e-9,
        'time': t.secs + t.nsecs*1e-9,
        'orient_x': msg.orientation.x,
        'orient_y': msg.orientation.y,
        'orient_z': msg.orientation.z,
        'orient_w': msg.orientation.w,
    }
    return data

class BagDataProcessor:
    def __init__(self, bag_path, topics, type_header):
        self.topics = topics
        self.bag_path = bag_path
        self.data_storage = {topic: [] for topic in topics}
        self.data_storage_np = {topic: None for topic in topics}
        self.topic_to_type = {topic: type_ for topic, type_ in zip(topics, type_header)}

    def read_topics_from_bag(self):
        bag = rosbag.Bag(self.bag_path)
        for topic, msg, t in bag.read_messages(topics=self.topics):
            if topic not in self.topics:
                continue
            if msg._type=='nav_msgs/Odometry':
                data = get_odomtry_data(t, msg)
            elif msg._type=='nlink_parser/LinktrackNodeframe3':
                data = get_nodeframe3_data(t, msg)
            elif msg._type=='sensor_msgs/Imu':
                data = get_imu_data(t, msg)
            else:
                continue
            self.data_storage[topic].append(data)

    def to_numpy(self):
        
        for topic, data in self.data_storage.items():
            self.data_storage_np[topic] = np.array(data)

    def save_to_file(self):
        """
        Save data to JSON files for each topic.
        """
        # Generate save directory
        bag_name = os.path.splitext(os.path.basename(self.bag_path))[0]
        output_dir = os.path.join(os.path.dirname(self.bag_path), bag_name)
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        # Save each topic's data to a JSON file
        for topic, data in self.data_storage.items():
            if not data:
                print("No data available for topic '{}'. Skipping save.".format(topic))
                continue

            # Construct output file path
            output_file = os.path.join(output_dir, "{}{}.json".format(self.topic_to_type[topic], topic.replace("/", "_")))

            # Save data to JSON
            with open(output_file, 'w') as json_file:
                json.dump(data, json_file, indent=4)
            print("Data for topic '{}' saved to '{}'.".format(topic, output_file))

